import json
from fastapi import FastAPI, HTTPException
from fastapi.websockets import WebSocket
import uvicorn
import json

app = FastAPI()

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    #获取websocket的ip,端口        
    client_info = websocket.client
    client_ip = client_info.host #['host']
    client_port = client_info.port # ['port']
    client_key = client_ip + ":" + str(client_port)
    client_node = None
    client_data = {"ip" : client_ip, "port" : client_port, 'key' : client_key }
    #ping_task = asyncio.create_task(self.ping_websocket(websocket))
    try:
        while True:
            #等待客户端发送消息
            try:
                data = await websocket.receive_text()
            except Exception as ex:
                print('websocket.receive_text异常')
                break

            print('WsServer收到消息:',data)

            inputs = json.loads(data)
            print(f"{inputs}")

            cmd=inputs.get('cmd','')
            if cmd=='register':
                print(f"in register, {inputs}")
                response_info = {"msg": "register ok", "code": 0, "data": client_data}
                response_info = json.dumps(response_info, ensure_ascii=False)
                await websocket.send_text( response_info )
            else:
                response_info = {"msg": "other cmd", "code": 0, "data": client_data}
                response_info = json.dumps(response_info, ensure_ascii=False)
                await websocket.send_text( response_info )

    except Exception as ex:
        print('WsServer客户端异常')
        print(ex)
    finally:
        print('WsServer客户端断开连接')
        await websocket.close()


# 启动服务器
if __name__ == "__main__":
    port=10086
    uvicorn.run(app, host="0.0.0.0", port=port)